#!/usr/bin/env python3

import sys
sys.path.append('../../../tools')
from gmcore_plot import *

parser = argparse.ArgumentParser(description='Plot test results')
parser.add_argument('-t', '--time-step', type=int, default=-1, help='Time step to plot')
parser.add_argument('-l', '--loop', action='store_true', help='Loop the plot')
parser.add_argument('-u', '--update-seconds', type=float, default=2, help='Update seconds')
parser.add_argument('-o', '--output', help='Output PDF figure')
args = parser.parse_args()

fig = plt.figure(figsize=(8, 8))
proj = ccrs.PlateCarree()

def plot(data_path, fig):
	f = addfile(data_path)

	time0 = parse_datetime(f.time[0])
	time1 = parse_datetime(f.time[args.time_step])
	fig.suptitle(f'Day {(time1 - time0).days}', fontsize=16)

	z_levels = np.linspace(8000, 10500, 13)

	gs = GridSpec(2, 2, figure=fig)

	ax = fig.add_subplot(gs[0,:], projection=proj)
	plot_contour_map(ax, f.gz[args.time_step,0,:,:] / g, cmap='jet', left_string=f'Geopotential height (m) at {((f.time[args.time_step] - f.time[0]).values).astype("timedelta64[D]")}', levels=z_levels)

	ax = fig.add_subplot(gs[1,0])
	plot_time_series(ax, f.time, f.te)

	ax = fig.add_subplot(gs[1,1])
	plot_time_series(ax, f.time, f.tpe)

plot('swm_rh.360x180.h0.nc', fig)
if args.loop:
	loop_plot('swm_rh.360x180.h0.nc', fig, plot, args.update_seconds)
elif args.output:
	plt.savefig(args.output, format='pdf')
else:
	plt.tight_layout()
	plt.show()
